import cv2
import numpy as np
import pyautogui
import win32gui
import win32con
import win32api
from typing import Optional, Tuple
import logging
import time
import os

class ChartCapturer:
    def __init__(self, account_id: str):
        self.account_id = account_id
        self.mt4_hwnd = None
        self.chart_handles = {}  # 缓存图表窗口句柄
        self.toolbar_height = 30  # 工具栏高度
        self.statusbar_height = 20  # 状态栏高度
        self.last_switch_time = None  # 添加上次切换时间记录
        
        # 创建调试图片目录
        try:
            # 获取当前工作目录（即运行脚本的目录）
            workspace_dir = os.getcwd()
            self.debug_dir = os.path.join(workspace_dir, "debug_images")
            print(f"调试图片目录路径: {self.debug_dir}")
            
            if not os.path.exists(self.debug_dir):
                try:
                    os.makedirs(self.debug_dir)
                    print(f"创建调试图片目录: {self.debug_dir}")
                except Exception as e:
                    print(f"创建调试图片目录失败: {str(e)}")
                    import traceback
                    traceback.print_exc()
            else:
                print(f"调试图片目录已存在: {self.debug_dir}")
        except Exception as e:
            print(f"设置调试图片目录失败: {str(e)}")
            import traceback
            traceback.print_exc()
        
    def set_mt4_window(self, hwnd: int) -> None:
        """设置MT4主窗口句柄"""
        self.mt4_hwnd = hwnd
        
    def capture_chart(self, symbol: str, timeframe: str, save_debug: bool = False) -> Optional[np.ndarray]:
        """截取指定货币对和时间周期的图表"""
        try:
            if not self.mt4_hwnd:
                print("未设置MT4窗口句柄")
                return None
                
            # 获取窗口标题（用于日志）
            title = win32gui.GetWindowText(self.mt4_hwnd)
            if not title:
                print("无法获取窗口标题")
                return None
            print(f"当前窗口标题: {title}")
                
            # 确保窗口可见且未最小化
            if not win32gui.IsWindowVisible(self.mt4_hwnd):
                print("MT4窗口不可见")
                return None
                
            # 获取图表区域
            chart_region = self.get_chart_region(self.mt4_hwnd)
            if not chart_region:
                print("无法获取图表区域")
                return None
                
            # 截取图表区域
            x, y, w, h = chart_region
            print(f"截取区域: x={x}, y={y}, width={w}, height={h}")
            
            # 确保区域有效
            if w <= 0 or h <= 0:
                print(f"无效的截图区域: width={w}, height={h}")
                return None
                
            screenshot = pyautogui.screenshot(region=(x, y, w, h))
            image = cv2.cvtColor(np.array(screenshot), cv2.COLOR_RGB2BGR)
            
            # 检查图像是否有效
            if image is None or image.size == 0:
                print("截图失败：图像为空")
                return None
                
            print(f"截图成功，图像尺寸: {image.shape}")
            
            # 保存原始截图
            if save_debug:
                debug_path = self._get_debug_image_path(symbol, timeframe, "debug")
                try:
                    cv2.imwrite(debug_path, image)
                    print(f"原始截图已保存: {debug_path}")
                except Exception as e:
                    print(f"保存截图失败: {str(e)}")
                    import traceback
                    traceback.print_exc()
            
            return image
            
        except Exception as e:
            print(f"截取图表失败: {str(e)}")
            import traceback
            traceback.print_exc()
            return None
        
    def _get_debug_image_path(self, symbol: str, timeframe: str, prefix: str = "debug") -> str:
        """
        生成调试图片的保存路径
        :param symbol: 交易对
        :param timeframe: 时间周期
        :param prefix: 文件名前缀
        :return: 完整的文件路径
        """
        filename = f"{prefix}_{symbol}_{timeframe}.png"
        path = os.path.join(self.debug_dir, filename)
        print(f"调试图片保存路径: {path}")
        return path
        
    def find_chart_window(self, symbol: str, timeframe: str) -> Optional[int]:
        """
        查找MT4图表窗口
        :param symbol: 交易对，例如"ETHUSD"
        :param timeframe: 时间周期，例如"H1"
        :return: 窗口句柄或None
        """
        try:
            # 先检查缓存
            symbol_timeframe = f"{symbol}_{timeframe}"
            if symbol_timeframe in self.chart_handles:
                hwnd = self.chart_handles[symbol_timeframe]
                if win32gui.IsWindow(hwnd) and win32gui.IsWindowVisible(hwnd):
                    print(f"使用缓存的窗口句柄: {hwnd}")
                    return hwnd
                else:
                    # 句柄无效，从缓存中移除
                    del self.chart_handles[symbol_timeframe]
            
            print(f"开始查找账号 {self.account_id} 的 {symbol} {timeframe} 图表窗口")
            
            # 查找窗口
            def callback(hwnd, extra):
                if win32gui.IsWindowVisible(hwnd):
                    title = win32gui.GetWindowText(hwnd)
                    print(f"找到窗口: {title}")
                    # 检查窗口标题是否包含账号和交易对信息
                    if (self.account_id in title and 
                        f"[{symbol},{timeframe}]" in title and 
                        "MetaTrader" not in title):  # 排除MetaEditor窗口
                        print(f"匹配到目标窗口，标题: {title}")
                        extra.append(hwnd)
                return True
                
            windows = []
            win32gui.EnumWindows(callback, windows)
            
            if windows:
                hwnd = windows[0]
                print(f"找到目标窗口，句柄: {hwnd}")
                
                # 恢复并激活窗口
                try:
                    # 如果窗口最小化，先恢复
                    if win32gui.IsIconic(hwnd):
                        print("窗口已最小化，正在恢复...")
                        win32gui.ShowWindow(hwnd, win32con.SW_RESTORE)
                        time.sleep(0.5)
                    
                    # 将窗口移到前台
                    print("正在激活窗口...")
                    win32gui.SetForegroundWindow(hwnd)
                    time.sleep(0.5)
                    
                    # 确保窗口可见
                    win32gui.ShowWindow(hwnd, win32con.SW_SHOW)
                    time.sleep(0.5)
                    
                    # 缓存窗口句柄
                    self.chart_handles[symbol_timeframe] = hwnd
                    
                except Exception as e:
                    print(f"激活窗口失败: {str(e)}")
                    
                return hwnd
                
            print("未找到目标窗口")
            return None
            
        except Exception as e:
            logging.error(f"查找MT4图表窗口失败: {str(e)}")
            return None
            
    def get_chart_region(self, hwnd: int) -> Tuple[int, int, int, int]:
        """
        获取图表区域
        :param hwnd: 窗口句柄
        :return: (left, top, width, height)
        """
        try:
            print("开始获取图表区域")
            # 获取窗口位置和大小
            left, top, right, bottom = win32gui.GetWindowRect(hwnd)
            print(f"窗口位置: left={left}, top={top}, right={right}, bottom={bottom}")
            
            # 检查窗口位置是否有效
            if left < -10000 or top < -10000:  # 窗口可能被最小化
                print("窗口位置无效，尝试恢复窗口...")
                win32gui.ShowWindow(hwnd, win32con.SW_RESTORE)
                time.sleep(0.5)
                left, top, right, bottom = win32gui.GetWindowRect(hwnd)
                print(f"恢复后窗口位置: left={left}, top={top}, right={right}, bottom={bottom}")
            
            # 获取客户区位置
            client_left, client_top, client_right, client_bottom = win32gui.GetClientRect(hwnd)
            client_left, client_top = win32gui.ClientToScreen(hwnd, (client_left, client_top))
            client_right, client_bottom = win32gui.ClientToScreen(hwnd, (client_right, client_bottom))
            print(f"客户区位置: left={client_left}, top={client_top}, right={client_right}, bottom={client_bottom}")
            
            # 计算图表区域，包含左侧交易对列表和上方信息
            chart_left = client_left  # 从客户区最左侧开始
            chart_top = client_top + self.toolbar_height  # 排除工具栏
            chart_width = client_right - client_left  # 使用整个宽度
            chart_height = client_bottom - chart_top - self.statusbar_height  # 排除状态栏
            
            # 确保区域有效
            if chart_width <= 0 or chart_height <= 0:
                print("警告：计算得到的图表区域无效，使用默认值")
                return (0, 0, 800, 600)
            
            region = (chart_left, chart_top, chart_width, chart_height)
            print(f"计算得到的图表区域: {region}")
            print(f"使用工具栏高度: {self.toolbar_height}, 状态栏高度: {self.statusbar_height}")
            return region
            
        except Exception as e:
            logging.error(f"获取图表区域失败: {str(e)}")
            return (0, 0, 800, 600)  # 默认区域
        
    def preprocess_image(self, image: np.ndarray, symbol: str, timeframe: str, save_debug: bool = False) -> np.ndarray:
        """预处理图表图像"""
        try:
            # 创建图像副本
            processed = image.copy()
            
            # 轻微增强对比度
            lab = cv2.cvtColor(processed, cv2.COLOR_BGR2LAB)
            l, a, b = cv2.split(lab)
            clahe = cv2.createCLAHE(clipLimit=1.5, tileGridSize=(8,8))
            cl = clahe.apply(l)
            enhanced = cv2.merge((cl,a,b))
            processed = cv2.cvtColor(enhanced, cv2.COLOR_LAB2BGR)
            
            # 保存处理后的图片
            if save_debug:
                processed_path = self._get_debug_image_path(symbol, timeframe, "processed")
                try:
                    cv2.imwrite(processed_path, processed)
                    print(f"处理后的图片已保存: {processed_path}")
                except Exception as e:
                    print(f"保存处理后的图片失败: {str(e)}")
                    import traceback
                    traceback.print_exc()
            
            return processed
            
        except Exception as e:
            logging.error(f"图像预处理失败: {str(e)}")
            return image
        
    def switch_chart(self) -> Optional[Tuple[str, str]]:
        """
        切换到下一个图表并返回当前图表信息
        :return: 当前图表的(货币对, 时间周期)元组，如果获取失败则返回None
        """
        try:
            # 记录切换时间
            self.last_switch_time = time.time()
            print(f"\n切换图表，记录时间戳: {self.last_switch_time}")
            
            # 找到MT4窗口
            if not self.mt4_hwnd:
                self.mt4_hwnd = self._find_mt4_window()
                if not self.mt4_hwnd:
                    print("未找到MT4窗口")
                    return None
            
            # 激活MT4窗口
            win32gui.ShowWindow(self.mt4_hwnd, win32con.SW_RESTORE)
            win32gui.SetForegroundWindow(self.mt4_hwnd)
            
            # 等待窗口激活
            time.sleep(0.5)
            
            # 使用Ctrl+Tab切换图表
            win32api.keybd_event(win32con.VK_CONTROL, 0, 0, 0)
            time.sleep(0.1)
            win32api.keybd_event(win32con.VK_TAB, 0, 0, 0)
            win32api.keybd_event(win32con.VK_TAB, 0, win32con.KEYEVENTF_KEYUP, 0)
            win32api.keybd_event(win32con.VK_CONTROL, 0, win32con.KEYEVENTF_KEYUP, 0)
            
            # 等待图表切换完成
            time.sleep(0.5)
            
            # 获取当前图表信息
            title = win32gui.GetWindowText(self.mt4_hwnd)
            if "[" in title and "]" in title:
                chart_info = title[title.find("[")+1:title.find("]")]
                if "," in chart_info:
                    symbol = chart_info.split(",")[0].strip()
                    timeframe = chart_info.split(",")[1].strip()
                    print(f"已切换到图表: {symbol}, {timeframe}")
                    return symbol, timeframe
                    
            print("无法从窗口标题解析图表信息")
            return None
            
        except Exception as e:
            logging.error(f"切换图表失败: {str(e)}")
            return None
            
    def _find_mt4_window(self) -> Optional[int]:
        """查找MT4主窗口"""
        try:
            def callback(hwnd, windows):
                if win32gui.IsWindowVisible(hwnd):
                    title = win32gui.GetWindowText(hwnd)
                    if "MetaTrader 4" in title:
                        windows.append(hwnd)
                return True
            
            windows = []
            win32gui.EnumWindows(callback, windows)
            
            if windows:
                return windows[0]
            return None
            
        except Exception as e:
            logging.error(f"查找MT4窗口失败: {str(e)}")
            return None
            
    def _get_chart_region(self) -> Optional[Tuple[int, int, int, int]]:
        """获取图表区域"""
        try:
            if not self.mt4_hwnd:
                return None
                
            # 获取窗口位置和大小
            left, top, right, bottom = win32gui.GetWindowRect(self.mt4_hwnd)
            
            # 获取客户区位置
            client_left, client_top, client_right, client_bottom = win32gui.GetClientRect(self.mt4_hwnd)
            client_left, client_top = win32gui.ClientToScreen(self.mt4_hwnd, (client_left, client_top))
            client_right, client_bottom = win32gui.ClientToScreen(self.mt4_hwnd, (client_right, client_bottom))
            
            # 计算图表区域（排除工具栏和状态栏）
            chart_top = client_top + 50  # 工具栏高度
            chart_bottom = client_bottom - 20  # 状态栏高度
            
            return (client_left, chart_top, client_right - client_left, chart_bottom - chart_top)
            
        except Exception as e:
            logging.error(f"获取图表区域失败: {str(e)}")
            return None
        
    def mark_kline(self, image: np.ndarray, kline_index: int, kline_width: int = 5, kline_gap: int = 1) -> np.ndarray:
        """
        在图表上标注指定的K线
        :param image: 原始图表图像
        :param kline_index: K线序号
        :param kline_width: K线宽度（像素）
        :param kline_gap: K线间距（像素）
        :return: 标注后的图像
        """
        try:
            # 创建图像副本
            marked_image = image.copy()
            
            # 计算K线的x坐标
            x = kline_index * (kline_width + kline_gap) + kline_width // 2
            
            # 在图像上画一条垂直线标注K线
            cv2.line(marked_image, (x, 0), (x, marked_image.shape[0]), (0, 255, 0), 2)
            
            # 在K线上方添加文字标注
            text = f"K线 {kline_index}"
            cv2.putText(marked_image, text, (x - 20, 20), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
            
            return marked_image
            
        except Exception as e:
            print(f"标注K线失败: {str(e)}")
            return image 